#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   train.py
@Time    :   2021-09-15 18:24:28
@Author  :   GuoLiuFang
@Version :   0.1
@Contact :   909104374@qq.com
@License :   (C)Copyright 2018-2021, RandomMatrix
@Desc    :   None
'''

# here put the import lib
from mmcv import Config
import json
import os.path as osp
import os
import sys
from dataset import build_data_loader
import torch

def get_args_parser(add_help=True):
    import argparse
    parser = argparse.ArgumentParser(description="Common PyTorch Training", add_help=True)
    parser.add_argument('config', help="config file path")
    parser.add_argument('--checkpoint', nargs='?', type=str, default=None)
    parser.add_argument('--resume', nargs='?', type=str, default=None)
    parser.add_argument('--device', default='1', help='GPU IDs', metavar='0 1 2 3 4 5 6 7')
    return parser

def main(args):
    # config 
    cfg = Config.fromfile(args.config)
    print(json.dumps(cfg._cfg_dict, indent=4))

    if args.checkpoint is not None:
        checkpoint_path = args.checkpoint
    else:
        cfg_name, _ = osp.splitext(osp.basename(args.config))
        checkpoint_path = osp.join('checkpoints', cfg_name)
    if not osp.isdir(checkpoint_path):
        os.makedirs(checkpoint_path)
    print('Checkpoint path: %s.' % checkpoint_path)
    sys.stdout.flush()

    # data loader
    data_loader = build_data_loader(cfg.data.train)
    train_loader = torch.utils.data.DataLoader(
        data_loader,
        batch_size=cfg.data.batch_size,
        shuffle=True,
        num_workers=8,
        drop_last=True,
        pin_memory=True
    )
    print("Train_Loader make success")
    pass

if __name__ == '__main__':
    args = get_args_parser().parse_args()
    main(args)